import pandas as pd
import os
import mplfinance as mpf
from tqdm import tqdm
import matplotlib.pyplot as plt

# --- 全局配置 ---
TRADES_LOG_PATH = 'all_trades_log.csv'
CACHE_DIR = 'feather_cache'

def visualize_trade(trade_info, stock_df, is_losing):
    """为单笔交易生成并保存带有标注的K线图 (稳定版)"""
    
    # 1. 准备数据
    buy_date = pd.to_datetime(trade_info['buy_datetime'])
    sell_date = pd.to_datetime(trade_info['sell_datetime'])
    symbol = trade_info['symbol']
    pnl_ratio = trade_info['pnl_ratio']

    # 2. 确定范围并修正日期
    start_plot_date = buy_date - pd.Timedelta(days=45)
    end_plot_date = sell_date + pd.Timedelta(days=15)
    
    # 扩大范围以确保RSI计算准确
    calc_start_date = buy_date - pd.Timedelta(days=80)
    calc_df = stock_df.loc[calc_start_date:end_plot_date].copy()
    if calc_df.empty: return

    # 3. 计算指标 (RSI 和均线)
    delta = calc_df['close'].diff()
    gain = delta.where(delta > 0, 0).ewm(span=14, adjust=False).mean()
    loss = -delta.where(delta < 0, 0).ewm(span=14, adjust=False).mean()
    rs = gain / loss
    calc_df['rsi'] = 100 - (100 / (1 + rs))
    
    plot_df = calc_df.loc[start_plot_date:].copy()
    if plot_df.empty: return

    buy_date = plot_df.index.asof(buy_date)
    sell_date = plot_df.index.asof(sell_date)
    if pd.isna(buy_date) or pd.isna(sell_date): return

    plot_df['SMA_5'] = plot_df['close'].rolling(5).mean()
    plot_df['SMA_10'] = plot_df['close'].rolling(10).mean()
    plot_df['SMA_20'] = plot_df['close'].rolling(20).mean()

    # 4. 获取买入时RSI并创建标记
    rsi_on_buy = plot_df.loc[buy_date, 'rsi'] if buy_date in plot_df.index else float('nan')
    rsi_str = f"{rsi_on_buy:.0f}" if pd.notna(rsi_on_buy) else "NaN"

    buy_markers = pd.Series(float('nan'), index=plot_df.index)
    sell_markers = pd.Series(float('nan'), index=plot_df.index)
    buy_markers.loc[buy_date] = plot_df.loc[buy_date, 'low'] * 0.98
    sell_markers.loc[sell_date] = plot_df.loc[sell_date, 'low'] * 0.98
    
    # 5. 准备绘图元素
    addplots = [
        mpf.make_addplot(buy_markers, type='scatter', marker='^', color='g', markersize=150),
        mpf.make_addplot(sell_markers, type='scatter', marker='v', color='r', markersize=150),
        mpf.make_addplot(plot_df[['SMA_5', 'SMA_10', 'SMA_20']])
    ]
    
    # 6. 绘图
    trade_type = "亏损" if is_losing else "盈利"
    charts_dir = os.path.join('charts', 'losing_trades' if is_losing else 'winning_trades')
    title = f"\n{trade_type}交易复盘: {symbol}\n买入: {buy_date.date()} (RSI: {rsi_str}) | 卖出: {sell_date.date()} | 盈亏: {pnl_ratio}"
    filename = os.path.join(charts_dir, f"{trade_type}_{symbol}_{buy_date.date()}_RSI{rsi_str}_{pnl_ratio.replace('%','')}.png")

    try:
        fig, _ = mpf.plot(
            plot_df,
            type='candle',
            style='charles',
            title=title,
            ylabel='价格 (Price)',
            volume=True,
            ylabel_lower='成交量 (Volume)',
            addplot=addplots,
            figsize=(16, 8),
            returnfig=True
        )
        fig.savefig(filename)
    except Exception as e:
        print(f"为 {symbol} 绘制图表时发生错误: {e}")
    finally:
        plt.close(fig) if 'fig' in locals() else None

def main():
    """主执行函数"""
    WINNING_CHARTS_DIR = 'charts/winning_trades'
    LOSING_CHARTS_DIR = 'charts/losing_trades'

    for path in [WINNING_CHARTS_DIR, LOSING_CHARTS_DIR]:
        if not os.path.exists(path):
            os.makedirs(path)
            print(f"图表输出目录已创建: {path}")

    if not os.path.exists(TRADES_LOG_PATH):
        print(f"错误: 交易日志 '{TRADES_LOG_PATH}' 不存在。请先运行 main.py 生成日志。")
        return

    trades_df = pd.read_csv(TRADES_LOG_PATH)

    if trades_df.empty:
        print("未发现任何交易，无需生成图表。")
        return
        
    print(f"共发现 {len(trades_df)} 笔交易，开始生成K线复盘图...")
    
    plt.switch_backend('agg')

    for _, trade in tqdm(trades_df.iterrows(), total=len(trades_df), desc="生成全部交易图表"):
        is_losing = trade['pnl'] < 0
        symbol = trade['symbol']
        cache_path = os.path.join(CACHE_DIR, f"{symbol}.feather")
        
        if not os.path.exists(cache_path):
            continue
            
        stock_df = pd.read_feather(cache_path)
        stock_df['date'] = pd.to_datetime(stock_df['date'])
        stock_df.set_index('date', inplace=True)
        
        visualize_trade(trade, stock_df, is_losing)
        
    print(f"\n所有交易的K线图已生成并保存到 'charts' 文件夹中。")

if __name__ == '__main__':
    main() 